'''
这个模块是用于提取图像特征的

'''

from PIL import Image
import os
from skimage.feature import hog
import joblib
import numpy as np
import variable as va

# 获取指定路径下的所有图片
def getImages(filePath, key):
    imgList = []
    label = va.labels[key]
    path = filePath + str(label)
    # 获取这个目录下的所有文件名
    fileList = os.listdir(path)
    for name in fileList:
        # 获取图片
        img = Image.open(os.path.join(path,name))
        # 改大小并加入列表（这个列表的格式是 [图片,文件名,类别标记]）
        imgList.append([img, name, key])
    return imgList


# 提取数据集中的图片的特征并保存
def get_feat(db_image, savePath):
    # 每次取出一种类别的图片
    i = 0
    for img in db_image:
        # 对图像进行一系列的处理
        img[0] = dealImage(img[0])
        # 获取特征
        fd = get_feat(img)
        fd = np.concatenate((fd, [img[2]]))
        fd_name = img[1] + '.feat'
        fd_path = os.path.join(savePath, fd_name)
        joblib.dump(fd, fd_path)
        i += 1

# 处理图像
def dealImage(img):
    # 将图片转化为特定的大小 width * height （最后一个参数是：抗锯齿）
    img = img.resize((va.image_width, va.image_height), Image.ANTIALIAS)
    # 将图片数组化
    img = np.reshape(img, (va.image_height, va.image_width, 3))    
    # 获取一张图片的灰度图（数组）
    img = img[:, :, 0] * 0.2989 + img[:, :, 1] * 0.5870 + img[:, :, 2] * 0.1140
    return img / 255.0

# 获取特征
def get_feat(img):
    # hog函数较难理解：
    # 参数分别是：1 - 要处理的图片
    # 2 - orientations 是bin的数量 ， 每个bin的角度是180/bin
    # 3 - block_norm ：指定norm的类型
    # 4 - pixels_per_cell : 每个cell的像素大小
    # 5 - cell_per_block ：每个块block中有多少个cell
    # 6 - visualise : 是否输梯度图
    feat = hog(img, orientations=12,block_norm='L1', pixels_per_cell=[8, 8],
               cells_per_block=[4, 4],visualize=False, transform_sqrt=True)
    return feat


# 获取文件夹中的原始数据集
def get_database(filePath):
    image_list = []
    # 依次获取一种类别的图片
    for key in va.labels.keys():
        # 获取数据集中的图片_这些图片的类别标记就是va.labels[key]
        images= getImages(filePath, key)
        #依次加到大列表中
        image_list += images

    # 返回这2个列表，分别代表这个数据集中所有的图片与它们的类别
    return image_list



# 获取特征，并保存在指定的文件夹下
def get_feats():
    # 加载训练集与测试集
    train_image = get_database(va.train_image_path)
    test_image = get_database(va.test_image_path)
    
    # 提取训练集与测试集的特征并保存
    get_feat(train_image, va.train_feat_path)
    get_feat(test_image, va.test_feat_path)
